import numpy as np
import math
from scipy.optimize import fsolve

from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt

class Surface:
    z = 0       # z coordinate
    radius = 0
    direction = True  #true:center lies right of surface, false:left of the surface
    nl = 0      # refractive index in the left
    nr = 0      # refractive index in the right
    aperture = 0

    def __init__(self, z, radius, dire, nl, nr, aper):
        self.z = z 
        self.radius = radius
        self.direction = dire
        self.nl = nl
        self.nr = nr
        self.aperture = aper
        return
    
    def plot(self):
        pts = []
        for i in range(10):
            for j in range(20):
                theta = 2*np.pi*j/20
                deltaphi = np.arcsin(self.aperture/self.radius)
                phi = i*deltaphi/10
                px = self.radius*math.sin(phi)*math.cos(theta)
                py = self.radius*math.sin(phi)*math.sin(theta)
                if self.direction == True:
                    pz = self.z+self.radius - self.radius*math.cos(phi)
                else:
                    pz = self.z-self.radius + self.radius*math.cos(phi)
                pts.append([px,py,pz])
        pts = np.array(pts)
        ax.plot_trisurf(pts[:,0],pts[:,1],pts[:,2],alpha = 0.8, color =[0.1,0.1,0.5])

class Lens:
    s1 = None  #surface 1
    s2 = None  #surface 2

    def __init__(self, r1, r2, z1, z2, n):
        self.s1 = Surface(z1, r1, True, 1.0, n)
        self.s2 = Surface(z2, r2, False, n, 1.0)
        return

class Light:
    hx = 0
    hy = 0
    z = 0
    kx = 0
    ky = 0
    kz = 0
    record_points = []

    def __init__(self, z, hx, hy, kx, ky):
        self.z = z
        self.hx = hx
        self.hy = hy
        self.kx = kx
        self.ky = ky
        self.kz = np.sqrt(1-kx*kx-ky*ky)
        self.record_points.append([self.z,self.hx,self.hy,self.kx,self.ky,self.kz])
        return
    
    def propagate(self,d):
        self.hx = self.hx + d*self.kx
        self.hy = self.hy + d*self.ky
        self.z = self.z + d*self.kz
        self.record_points.append([self.z,self.hx,self.hy,self.kx,self.ky,self.kz])
        return

    def f_left(self,x,surf):
        h1 = np.sqrt(math.pow(surf.radius,2)-math.pow(surf.radius-x,2))
        h2x = self.hx + (x+surf.z-self.z)*(self.kx/self.kz)
        h2y = self.hy + (x+surf.z-self.z)*(self.ky/self.kz)
        h2 = np.sqrt(math.pow(h2x,2)+math.pow(h2y,2))
        return h1-h2
    
    def f_right(self,x,surf):
        h1 = np.sqrt(math.pow(surf.radius,2)-math.pow(surf.radius-np.abs(x),2))
        h2x = self.hx + (surf.z-self.z-x)*(self.kx/self.kz)
        h2y = self.hy + (surf.z-self.z-x)*(self.ky/self.kz)
        h2 = np.sqrt(math.pow(h2x,2)+math.pow(h2y,2))
        return h1-h2

    def next_surface(self, surf):
        if surf.direction == True:
            sol = fsolve(self.f_left, 0,args = (surf))
            dz = sol[0]+(surf.z-self.z)
        elif surf.direction == False:
            sol = fsolve(self.f_right, 0,args = (surf))
            dz = (surf.z-self.z)-sol[0]
        else:
            print("Error in finding next_surface!!!")
        
        self.propagate(dz/self.kz)
        return
    
    def next_surface_analytic(self, surf):
        if surf.direction == True:
            vcent = np.array([-self.hx,-self.hy,(surf.z+surf.radius)-self.z])
        elif surf.direction == False:
            vcent = np.array([-self.hx,-self.hy,(surf.z-surf.radius)-self.z])
        prop_direct = np.array([self.kx,self.ky,self.kz])
        d1 = np.dot(vcent, prop_direct)
        distance_sq = np.dot(vcent,vcent) - d1*d1
        inside = np.dot(vcent,vcent) - surf.radius**2
        d = d1 - np.sqrt(surf.radius**2 - distance_sq)*np.sign(inside)
        self.propagate(np.abs(d)) 
        
    def refract(self, surf):
        if(surf.direction == True):
            normal_k = np.array([self.hx,self.hy,-surf.radius])
            normal_k = normal_k/np.linalg.norm(normal_k)
            incident_k = np.array([self.kx,self.ky,self.kz])
            incident_k = incident_k/np.linalg.norm(incident_k)
            cosine_theta_i = np.dot(normal_k, incident_k)
            sine_theta_i = np.sqrt(1-math.pow(cosine_theta_i,2))
            sine_theta_r = sine_theta_i*surf.nl/surf.nr
            cosine_theta_r = np.sign(cosine_theta_i)*np.sqrt(1-math.pow(sine_theta_r,2))
            kr_vert = sine_theta_r*(incident_k-cosine_theta_i*normal_k)/sine_theta_i
            kr_perp = cosine_theta_r*normal_k
            kr = kr_vert + kr_perp

        elif(surf.direction == False):
            normal_k = np.array([self.hx,self.hy,surf.radius])
            normal_k = normal_k/np.linalg.norm(normal_k)
            incident_k = np.array([self.kx,self.ky,self.kz])
            incident_k = incident_k/np.linalg.norm(incident_k)
            cosine_theta_i = np.dot(normal_k, incident_k)
            sine_theta_i = np.sqrt(1-math.pow(cosine_theta_i,2))
            sine_theta_r = sine_theta_i*surf.nl/surf.nr
            cosine_theta_r = np.sign(cosine_theta_i)*np.sqrt(1-math.pow(sine_theta_r,2))
            kr_vert = sine_theta_r*(incident_k-cosine_theta_i*normal_k)/sine_theta_i
            kr_perp = cosine_theta_r*normal_k
            kr = kr_vert + kr_perp

        else:
            print("Error in refracting!!!")

        self.kx = kr[0]
        self.ky = kr[1]
        self.kz = kr[2]
        return

    def plot(self):
        pts = np.array(self.record_points)
        ax.plot3D(pts[:,1],pts[:,2],pts[:,0]) 

if __name__ == "__main__":
    fig = plt.figure()
    ax = plt.axes(projection = '3d')

    s1 = Surface(200,100, True, 1.0, 1.5, 40)
    s2 = Surface(220,100, False, 1.5, 1.0, 40)
    for i in range(10):
        l = Light(0, -5, 0, i*np.pi/180,2*np.pi/180)
        l.next_surface_analytic(s1)
        l.refract(s1)
        l.next_surface_analytic(s2)
        l.refract(s2)
        l.propagate(100)
        l.plot()
    
    s1.plot()
    s2.plot()
    
    plt.show()
